/*-
 * Copyright (C) 2018-2020, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>, 赵楷文 (Kaiwen
 * Zhao) <zhaokaiwen@ticpsh.com>
 */

#include "testvisitor.h"

#include "common/AttrAnalyzer.h"
#include "common/Coder.h"

using namespace std;
using namespace clang;
using namespace rocket::common;
namespace fs = boost::filesystem;

/********************************************************************************/
/* TestTopLevelVisitor 提取待测函数和全局变量 */

void rocket::testif::TestTopLevelVisitor::addOverriding(
    const CXXMethodDecl *overridden, const CXXMethodDecl *overridding) {
  auto method_it = FuncDeclAction(overridden);
  // if (overridding->isTemplated()) return;
  method_it->second.overridings.insert(overridding);
  for (auto over : overridden->overridden_methods()) {
    addOverriding(over, overridding);
  }
}

std::map<const clang::FunctionDecl *, rocket::testif::FuncDeclInfo>::iterator
rocket::testif::TestTopLevelVisitor::FuncDeclAction(const FunctionDecl *fd) {
  fd = fd->getFirstDecl();
  auto func_it = functiondecls_.find(fd);
  if (func_it == functiondecls_.end()) {
    bool ret;
    tie(func_it, ret) = functiondecls_.insert(make_pair(fd, FuncDeclInfo()));
    // assert(ret);

    if (auto md = dyn_cast<CXXMethodDecl>(fd)) {
      for (auto overridden : md->overridden_methods()) {
        addOverriding(overridden, md);
      }
    }
  }
  return func_it;
}

/**
 * @brief 筛选目标待测函数
 * @return
 */
bool rocket::testif::TestTopLevelVisitor::VisitFunctionDecl(FunctionDecl *fd) {
  if (!fd) return true;

  auto fst_fd = fd->getFirstDecl();
  if (fst_fd) FuncDeclAction(fst_fd);

  if (fd->getASTContext().getSourceManager().isInMainFile(fd->getLocation()) ||
      analyze_level == AnalyzeLevel::AL_Full) {
    addTargetFunction(fd);
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::VisitFunctionTemplateDecl(
    FunctionTemplateDecl *ftd) {
  if (!ftd) return true;

  {
    auto fst_ftd = ftd->getTemplatedDecl()->getFirstDecl();
    if (fst_ftd) {
      FuncDeclAction(fst_ftd);
    }

    for (auto ftd_spec : ftd->specializations()) {
      auto fst_spec = ftd_spec->getFirstDecl();
      if (fst_spec) {
        FuncDeclAction(fst_spec);
        VisitFunctionDecl(fst_spec);
      }
    }
  }

  if (ftd && (ftd->getASTContext().getSourceManager().isInMainFile(
                  ftd->getLocation()) ||
              analyze_level == AnalyzeLevel::AL_Full)) {
    addTargetFunction(ftd);
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::VisitVarDecl(VarDecl *vd) {
  try {
    if (isLocalGlobal(vd)) {
      if (!vd->isFirstDecl())
        vd = vd->getFirstDecl();  // 类内static的getDefination会产生空指针
      if (globals_.find(vd) == globals_.end()) {
        globals_.insert(make_pair(vd, TargetGlobalVar(vd)));
      }
    }
  } catch (...) {
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::VisitClassTemplateDecl(
    clang::ClassTemplateDecl *ctd) {
  // FIXED: 复杂项目中直接访问类特化会导致因数据量过多从而段错误
  // NOTE: 非全解析时，由于挑选机制调整，模板特化可通过被其他函数调用而被检索到
  if (analyze_level != AnalyzeLevel::AL_Full) return true;

  for (auto ctsd : ctd->specializations()) {
    for (auto stub_decl : ctsd->decls()) {
      share::FixedAddressVisitor<TestTopLevelVisitor>::TraverseDecl(stub_decl);
    }
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::VisitEnumDecl(clang::EnumDecl *ed) {
  if (!ed) return true;
  if (auto edt = ed->getTypeForDecl()) {
    types.insert(edt);
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::TraverseDecl(clang::Decl *dl) {
  RecursiveASTVisitor<TestTopLevelVisitor>::TraverseDecl(dl);
  return true;
}

bool rocket::testif::TestTopLevelVisitor::addTargetFunction(FunctionDecl *fd) {
  if (!fd) return false;

  if (functions_.count(fd)) return false;

  if (!fd->doesThisDeclarationHaveABody()) {
    auto fd_d = fd->getDefinition();
    if (fd_d != nullptr && fd_d != fd) return addTargetFunction(fd_d);

    auto fd_o =
        fd->getMemberSpecializationInfo()
            ? dyn_cast<FunctionDecl>(
                  fd->getMemberSpecializationInfo()->getInstantiatedFrom())
            : nullptr;
    if (!fd_o || !fd_o->getDefinition()) {
      return false;
    }
  }

  if (CXXMethodDecl *cxxmd = dyn_cast<CXXMethodDecl>(fd)) {
    // 由编译器生成的函数不在待测函数中
    if (cxxmd->isDefaulted() || cxxmd->isDeleted()) return false;
  }

  auto tf_it =
      functions_.insert(make_pair(fd->getFirstDecl(), TargetFunction(fd)));
  if (!tf_it.second) return false;
  TargetFunction &tf = tf_it.first->second;
  analyzeTargetFunction(tf);
  if (fd->isTemplateInstantiation()) {
    if (auto temp_info = fd->getTemplateSpecializationInfo()) {
      //      functions_.insert(make_pair(fd->getFirstDecl(), tf));
      auto temp_orign =
          temp_info->getTemplate()->getTemplatedDecl()->getFirstDecl();
      auto temp_it = functions_.find(temp_orign);
      if (temp_it == functions_.end()) {
        auto temp_orign_it = functions_.insert(
            make_pair(temp_orign, TargetFunction(temp_orign)));
        if (!temp_orign_it.second) return false;
        TargetFunction &to = temp_orign_it.first->second;
        analyzeTargetFunction(to);
        to.specfs.insert(make_pair(fd->getFirstDecl(), tf));
      } else {
        temp_it->second.specfs.insert(make_pair(fd->getFirstDecl(), tf));
      }
      //      return true;
    }
  }

  {
    FunctionDecl *func_orign = nullptr;
    FunctionTemplateDecl *func_template = nullptr;

    if (fd->isTemplateDecl()) {
      func_template = fd->getDescribedFunctionTemplate();
    } else if (fd->isFunctionTemplateSpecialization()) {
      func_orign = fd->getPrimaryTemplate()->getTemplatedDecl();
      func_template = func_orign->getDescribedFunctionTemplate();
    }

    // func_orign
    if (fd->getMemberSpecializationInfo()) {
      func_orign = llvm::dyn_cast<clang::FunctionDecl>(
          fd->getMemberSpecializationInfo()->getInstantiatedFrom());
    } else if (func_template) {
      if (auto ift = func_template->getInstantiatedFromMemberTemplate()) {
        func_orign = ift->getTemplatedDecl();
      }
    }

    addTargetFunction(func_orign);
    addTargetFunction(func_template);
  }

  //  functions_.insert(make_pair(fd->getFirstDecl(), tf));
  return true;
}

bool rocket::testif::TestTopLevelVisitor::addTargetFunction(
    clang::FunctionTemplateDecl *ftd) {
  if (!ftd) return false;

  if (auto ift = ftd->getInstantiatedFromMemberTemplate()) {
    if (!ift->getTemplatedDecl()->getDefinition()) {
      return true;
    }
  } else if (!ftd->getTemplatedDecl()->doesThisDeclarationHaveABody()) {
    return true;
  }

  auto td = ftd->getTemplatedDecl()->getFirstDecl();
  auto f_it = functions_.find(td);
  if (f_it != functions_.end()) {
    // (*f_it).second = tf;
  } else {
    TargetFunction tf(ftd);
    analyzeTargetFunction(tf);
    functions_.insert(make_pair(td, tf));
  }
  return true;
}

bool rocket::testif::TestTopLevelVisitor::analyzeTargetFunction(
    TargetFunction &tf) {
  TestIfVisitor tiv(tf.declptr->getASTContext(), tf, setting);
  tiv.TraverseDecl(tf.declptr);

  for (auto param : tf.declptr->parameters()) {
    tiv.TraverseDecl(param);
    if (ast.getLangOpts().CPlusPlus) {
      if (auto cxxrd = param->getType()->getAsCXXRecordDecl()) {
        for (auto ctor : cxxrd->ctors()) {
          if (ctor->isCopyConstructor() && !ctor->isDeleted()) {
            tf.funcs.insert(ctor->getFirstDecl());
          }
        }
      }
    }
  }

  if (auto ctor = llvm::dyn_cast<CXXConstructorDecl>(tf.declptr)) {
    if (setting.analyze_indirect_stubs && ctor) {
      for (auto ctorinit : ctor->inits()) {
        if (auto init = ctorinit->getInit()) {
          tiv.TraverseStmt(init);
        }
      }

      if (auto cxxrd = ctor->getParent()) {
        for (auto field : cxxrd->fields()) {
          if (auto field_init = field->getInClassInitializer()) {
            tiv.TraverseStmt(field_init);
          }
        }
      }
    }
  }

  tf.fixed_addresses = tiv.fixed_addresses;

  set<FunctionDecl *> stub_functions;
  for (CallExpr *ce : tf.stubs) {
    clang::FunctionDecl *call_decl = ce->getDirectCallee();
    if (call_decl) {
      if (call_decl->getNameAsString().compare(0, sizeof("__builtin") - 1,
                                               "__builtin") == 0) {
        continue;  // __builtin函数
      } else if (auto method = dyn_cast<CXXMethodDecl>(call_decl)) {
        if (method->getParent() ==
            tf.declptr->getASTContext().getVaListTagDecl())
          continue;
      }

      if (!call_decl->getType()->getAs<FunctionNoProtoType>()) {
        stub_functions.insert(call_decl);
      }
    }
  }
  for (clang::FunctionDecl *func_decl : tf.funcs) {
    if (func_decl->getNameAsString().compare(0, sizeof("__builtin") - 1,
                                             "__builtin") == 0) {
      continue;  // __builtin函数
    } else if (auto method = dyn_cast<CXXMethodDecl>(func_decl)) {
      if (method->getParent() == tf.declptr->getASTContext().getVaListTagDecl())
        continue;
    }
    stub_functions.insert(func_decl->getFirstDecl());
  }
  for (const CXXConstructExpr *ctor : tf.ctors) {
    auto cxxrd = ctor->getConstructor()->getParent();
    if (cxxrd == tf.declptr->getASTContext().getVaListTagDecl()) {
      continue;
    }
    stub_functions.insert(ctor->getConstructor() /*->getFirstDecl()*/);
  }

  if (analyze_level != AnalyzeLevel::AL_Local) {
    for (auto stub_decl : stub_functions) {
      addTargetFunction(stub_decl);
    }
  }

  for (auto &tf_spec : tf.specfs) {
    analyzeTargetFunction(tf_spec.second);
  }

  return true;
}

/********************************************************************************/
/* InnerGlobalVisitor 提取待某一结构中使用的全局变量*/

// InnerGlobalVisitor
bool rocket::testif::InnerGlobalVisitor::VisitDeclRefExpr(DeclRefExpr *dre) {
  Decl *decl = dre->getDecl();
  if (VarDecl *vd = dyn_cast<VarDecl>(decl)) {
    /*if (isLocalGlobal(vd))*/ { globals_.insert(vd->getFirstDecl()); }
  } else if (FunctionDecl *func = dyn_cast<FunctionDecl>(decl)) {
    functions_.insert(func /*->getFirstDecl()*/);
  }
  return true;
}

bool rocket::testif::InnerGlobalVisitor::VisitCXXConstructExpr(
    CXXConstructExpr *ctor) {
  if (auto func = ctor->getConstructor()) {
    functions_.insert(func);
  }
  return true;
}

bool rocket::testif::InnerGlobalVisitor::VisitUserDefinedLiteral(
    UserDefinedLiteral *udl) {
  if (udl) udliteral_.insert(udl);
  return true;
}

bool rocket::testif::InnerGlobalVisitor::VisitMemberExpr(
    clang::MemberExpr *me) {
  member_exprs.push_back(me);
  return true;
}

/********************************************************************************/
/* TestIfVisitor 解析待测函数*/

string rocket::testif::getScopeTypeAsString(
    const rocket::testif::SCOPEKIND &st) {
  return SCOPEKIND_STR[st];
}

bool rocket::testif::TestIfVisitor::dealUntraversedNode(clang::Expr *expr) {
  if (traversedExpr.find(expr) != traversedExpr.end()) {
    return false;
  }

  if (auto initList = dyn_cast<InitListExpr>(expr)) {
    for (auto child : initList->children()) {
      if (auto expr = dyn_cast<Expr>(child)) {
        dealUntraversedNode(expr);
      }
    }
  } else if (auto tExpr = dyn_cast<MaterializeTemporaryExpr>(expr)) {
    VisitMaterializeTemporaryExpr(tExpr);
  } else if (auto ce = dyn_cast<CallExpr>(expr)) {
    VisitCallExpr(ce);
  } else if (auto dre = dyn_cast<DeclRefExpr>(expr)) {
    VisitDeclRefExpr(dre);
  } else if (auto cxxce = dyn_cast<CXXConstructExpr>(expr)) {
    VisitCXXConstructExpr(cxxce);
  } else if (auto ce = dyn_cast<CastExpr>(expr)) {
    VisitCastExpr(ce);
  } else if (auto te = dyn_cast<CXXThrowExpr>(expr)) {
    VisitCXXThrowExpr(te);
  } else if (auto me = dyn_cast<MemberExpr>(expr)) {
    VisitMemberExpr(me);
  } else if (auto cxxdae = dyn_cast<CXXDefaultArgExpr>(expr)) {
    VisitCXXDefaultArgExpr(cxxdae);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitMaterializeTemporaryExpr(
    clang::MaterializeTemporaryExpr *tExpr) {
  for (auto child : tExpr->children()) {
    if (auto expr = dyn_cast<Expr>(child)) {
      dealUntraversedNode(expr);
    }
  }
  return true;
}

/**
 * @brief 用于对目标函数中桩函数的额外解析
 * @param tf
 * @param fd
 * @return
 */
bool rocket::testif::TestIfVisitor::analyzeStubFunction(
    clang::FunctionDecl *fd) {
  if (!fd) return false;
  auto ret_type = fd->getReturnType();
  if (!ret_type->isRecordType()) return false;
  auto cxxrd = ret_type->getAsCXXRecordDecl();
  if (!cxxrd) return false;

  for (auto ctor : cxxrd->ctors()) {
    if (ctor->isCopyConstructor()) {
      func_.funcs.insert(ctor->getFirstDecl());
    }
  }

  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXNewExpr(clang::CXXNewExpr *ne) {
  if (auto new_opt = ne->getOperatorNew()) {
    if (!new_opt->isImplicit()) func_.funcs.insert(new_opt);
  }
  if (auto delete_opt = ne->getOperatorDelete()) {
    if (!delete_opt->isImplicit()) func_.funcs.insert(delete_opt);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXDeleteExpr(
    clang::CXXDeleteExpr *de) {
  if (auto delete_opt = de->getOperatorDelete()) {
    if (!delete_opt->isImplicit()) func_.funcs.insert(delete_opt);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitDeclStmt(clang::DeclStmt *ds) {
  if (!ds) return true;

  for (auto decl : ds->decls()) {
    if (auto vd = llvm::dyn_cast<VarDecl>(decl)) {
      if (vd->isLocalVarDecl() && vd->getStorageClass() != clang::SC_Extern &&
          vd->getStorageClass() != clang::SC_Static) {
        func_.locals_decls.insert(make_pair(vd, ds));
      }
    }
  }

  return true;
}

bool rocket::testif::TestIfVisitor::VisitCallExpr(CallExpr *ce) {
  func_.stubs.insert(ce);
  analyzeStubFunction(ce->getDirectCallee());
  return true;
}

bool rocket::testif::TestIfVisitor::VisitDeclRefExpr(DeclRefExpr *dre) {
  Decl *decl = dre->getDecl();
  if (VarDecl *vd = dyn_cast<VarDecl>(decl)) {
    if (isLocalGlobal(vd)) {
      auto &globals = func_.globals;
      globals.insert(vd);
      // if (auto init = vd->getInit()) {
      //   if (traversedExpr.find(init) == traversedExpr.end()) {
      //     traversedExpr.insert(init);
      //     TraverseStmt(init);
      //   }
      // }
    }
  } else if (FunctionDecl *fd = dyn_cast<FunctionDecl>(decl)) {
    auto &funcs = func_.funcs;
    analyzeStubFunction(fd);
    funcs.insert(fd);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXConstructExpr(
    CXXConstructExpr *cxxce) {
  func_.ctors.insert(cxxce);
  auto ctor = cxxce->getConstructor();
  if (setting.analyze_indirect_stubs && ctor &&
      ctor->isTemplateInstantiation()) {
    for (auto param : ctor->parameters()) {
      if (traversedDecl.find(param) == traversedDecl.end()) {
        traversedDecl.insert(param);
        VisitDecl(param);
      }
    }
    for (auto ctorinit : ctor->inits()) {
      if (auto init = ctorinit->getInit()) {
        if (traversedExpr.find(init) == traversedExpr.end()) {
          traversedExpr.insert(init);
          TraverseStmt(init);
        }
      }
    }

    if (auto body = ctor->getBody()) {
      if (traversedStmt.find(body) == traversedStmt.end()) {
        traversedStmt.insert(body);
        TraverseStmt(body);
      }
    }
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitDecl(Decl *decl) {
  if (VarDecl *vd = clang::dyn_cast<VarDecl>(decl)) {
    if (vd->isStaticLocal()) {
      func_.statics.push_back(vd);
    }
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitCastExpr(clang::CastExpr *expr) {
  func_.types.emplace_back(expr->getType());
  if (auto sub_expr = expr->getSubExpr()) {
    auto sub_type = sub_expr->getType();
    if (!sub_type.isNull()) func_.types.emplace_back(sub_type);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXThrowExpr(CXXThrowExpr *te) {
  func_.throwexprs.push_back(te);
  return true;
}

bool rocket::testif::TestIfVisitor::VisitMemberExpr(clang::MemberExpr *me) {
  func_.memberexprs.push_back(me);
  auto mem_decl = me->getMemberDecl();
  if (auto mem_static_decl = dyn_cast<VarDecl>(mem_decl)) {
    func_.globals.insert(mem_static_decl);
  }
  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXForRangeStmt(
    clang::CXXForRangeStmt *stmt) {
  if (auto range_stmt = dyn_cast<clang::CXXForRangeStmt>(stmt)) {
    if (auto rinit = range_stmt->getRangeInit())
      clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(rinit);

    if (auto begin_stmt = range_stmt->getBeginStmt()) {
      for (auto begin_decl : begin_stmt->decls()) {
        if (auto vd = dyn_cast<VarDecl>(begin_decl)) {
          if (auto init = vd->getInit()) {
            clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(init);
          }
        }
      }
    }

    if (auto end_stmt = range_stmt->getEndStmt()) {
      for (auto end_decl : end_stmt->decls()) {
        if (auto vd = dyn_cast<VarDecl>(end_decl)) {
          if (auto init = vd->getInit()) {
            clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(init);
          }
        }
      }
    }

    if (auto cond = range_stmt->getCond())
      clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(cond);

    if (auto inc = range_stmt->getInc())
      clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(inc);

    if (auto lv_stmt = range_stmt->getLoopVarStmt()) {
      for (auto loopvar_decl : lv_stmt->decls()) {
        if (auto vd = dyn_cast<VarDecl>(loopvar_decl)) {
          if (auto init = vd->getInit()) {
            clang::RecursiveASTVisitor<TestIfVisitor>::TraverseStmt(init);
          }
        }
      }
    }

    return true;
  }

  return true;
}

bool rocket::testif::TestIfVisitor::VisitCXXDefaultArgExpr(
    clang::CXXDefaultArgExpr *expr) {
  if (!expr) return true;
  auto param_decl = expr->getParam();
  if (!param_decl) return true;
  auto param_default = param_decl->getDefaultArg();

  if (param_default)
    TestIfVisitor::TraverseStmt(expr->getParam()->getDefaultArg());

  return true;
}

rocket::testif::TargetFunction::TargetFunction(
    clang::FunctionTemplateDecl *decl_)
    : declptr(decl_->getTemplatedDecl()) {
  for (auto specf : decl_->specializations()) {
    specfs.insert(std::make_pair(specf, TargetFunction(specf)));
  }
}
